package com.example.scribe.backend.handlers;

import com.example.scribe.SyntaxNode;
import com.example.scribe.backend.ASTUtils;
import com.example.scribe.backend.BytecodeOps;
import com.example.scribe.backend.ExpressionEmitter;
import com.example.scribe.backend.ExpressionNodeHandler;
import com.example.scribe.backend.LocalVariableAllocator;
import net.bytebuddy.implementation.bytecode.StackManipulation;
import net.bytebuddy.implementation.bytecode.StackManipulation.Compound;

import java.lang.reflect.Method;

/**
 * Handler for range expressions (0..10 or 0..=10).
 * Generates bytecode that returns an Iterable<Integer>.
 */
public final class RangeExpressionHandler implements ExpressionNodeHandler {
    private final ExpressionEmitter exprEmitter;

    public RangeExpressionHandler(ExpressionEmitter exprEmitter) {
        this.exprEmitter = exprEmitter;
    }

    @Override
    public boolean canHandle(String nodeType) {
        return "RangeExpr".equals(nodeType);
    }

    @Override
    public StackManipulation handle(SyntaxNode node, String generatedClassName, LocalVariableAllocator locals) {
        if (node == null) {
            return StackManipulation.Trivial.INSTANCE;
        }

        // RangeExpr structure: Start (with expression), Operator, End (with expression)
        SyntaxNode startWrapper = null;
        SyntaxNode endWrapper = null;
        String operator = null;

        for (SyntaxNode child : node.children()) {
            if ("Start".equals(child.type())) {
                startWrapper = child;
            } else if ("End".equals(child.type())) {
                endWrapper = child;
            } else if ("Operator".equals(child.type())) {
                operator = child.text();
            }
        }

        if (startWrapper == null || endWrapper == null || operator == null) {
            return StackManipulation.Trivial.INSTANCE;
        }

        // Extract actual expression nodes from Start and End wrappers
        SyntaxNode startExpr = null;
        SyntaxNode endExpr = null;
        
        // Start and End wrappers should have the actual expression as their first child
        if (!startWrapper.children().isEmpty()) {
            startExpr = startWrapper.children().get(0);
        }
        if (!endWrapper.children().isEmpty()) {
            endExpr = endWrapper.children().get(0);
        }
        
        // Fallback: try to find expression nodes recursively
        if (startExpr == null) {
            startExpr = findExpressionNode(startWrapper);
        }
        if (endExpr == null) {
            endExpr = findExpressionNode(endWrapper);
        }

        if (startExpr == null || endExpr == null) {
            return StackManipulation.Trivial.INSTANCE;
        }

        // Emit start and end expressions
        StackManipulation startEmit = exprEmitter.emit(startExpr, generatedClassName, locals);
        StackManipulation endEmit = exprEmitter.emit(endExpr, generatedClassName, locals);

        if (startEmit == StackManipulation.Trivial.INSTANCE || endEmit == StackManipulation.Trivial.INSTANCE) {
            // Try to parse as numbers directly from expression nodes
            if (startEmit == StackManipulation.Trivial.INSTANCE && ASTUtils.isNumber(startExpr)) {
                try {
                    int startVal = Integer.parseInt(startExpr.text());
                    startEmit = BytecodeOps.pushInt(startVal);
                } catch (NumberFormatException e) {
                    return StackManipulation.Trivial.INSTANCE;
                }
            }
            if (endEmit == StackManipulation.Trivial.INSTANCE && ASTUtils.isNumber(endExpr)) {
                try {
                    int endVal = Integer.parseInt(endExpr.text());
                    endEmit = BytecodeOps.pushInt(endVal);
                } catch (NumberFormatException e) {
                    return StackManipulation.Trivial.INSTANCE;
                }
            }
            if (startEmit == StackManipulation.Trivial.INSTANCE || endEmit == StackManipulation.Trivial.INSTANCE) {
                return StackManipulation.Trivial.INSTANCE;
            }
        }

        // Determine which method to call based on operator
        Method rangeMethod;
        try {
            if ("..".equals(operator)) {
                // Exclusive range: 0..10 means 0 to 9
                rangeMethod = com.example.scribe.backend.RangeHelper.class.getMethod("range", int.class, int.class);
            } else if ("..=".equals(operator)) {
                // Inclusive range: 0..=10 means 0 to 10
                rangeMethod = com.example.scribe.backend.RangeHelper.class.getMethod("rangeInclusive", int.class, int.class);
            } else {
                return StackManipulation.Trivial.INSTANCE;
            }
        } catch (NoSuchMethodException e) {
            return StackManipulation.Trivial.INSTANCE;
        }

        // Call RangeHelper.range(start, end) or RangeHelper.rangeInclusive(start, end)
        // Returns Iterable<Integer>
        StackManipulation invokeRange = BytecodeOps.invokeStatic(rangeMethod);

        // Stack: [start, end] -> [Iterable<Integer>]
        return new Compound(startEmit, endEmit, invokeRange);
    }

    /**
     * Find the actual expression node from a Start or End node.
     * The Start/End nodes contain text, but we need to find the actual expression AST node.
     */
    private SyntaxNode findExpressionNode(SyntaxNode node) {
        if (node == null) {
            return null;
        }

        // If the node has children, look for expression nodes
        if (!node.children().isEmpty()) {
            // Look for expression, assignment, conditional, logical_or, etc.
            for (SyntaxNode child : node.children()) {
                String type = child.type();
                if ("expression".equals(type) || "assignment".equals(type) || 
                    "conditional".equals(type) || "logical_or".equals(type) ||
                    "logical_and".equals(type) || "equality".equals(type) ||
                    "comparison".equals(type) || "bitwise_or".equals(type) ||
                    "bitwise_xor".equals(type) || "bitwise_and".equals(type) ||
                    "addition".equals(type) || "subtraction".equals(type) ||
                    "multiplication".equals(type) || "shift".equals(type) ||
                    "unary".equals(type) || "call".equals(type)) {
                    return child;
                }
                // Also check for primary nodes with identifier or number
                if ("primary".equals(type)) {
                    return child;
                }
            }
            // If no structured node found, recursively search
            for (SyntaxNode child : node.children()) {
                SyntaxNode found = findExpressionNode(child);
                if (found != null) {
                    return found;
                }
            }
        }

        // If node itself is an identifier or number, return it
        if (ASTUtils.isIdentifier(node) || ASTUtils.isNumber(node)) {
            return node;
        }

        return null;
    }
}

